#include <mcp/auth/authenticator.h>
#include <mcp/error.h>
#include <mcp/json/value.h>

#include <iostream>
#include <sstream>
#include <chrono>
#include <random>
#include <memory>
#include <unordered_map>
#include <mutex>

namespace mcp {
namespace auth {

// Helper functions
static std::string base64_encode(const std::string& data) {
    static const std::string chars =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

    std::string result;
    result.reserve(((data.size() + 2) / 3) * 4);

    for (size_t i = 0; i < data.size(); i += 3) {
        uint32_t value = 0;
        for (size_t j = 0; j < 3 && i + j < data.size(); ++j) {
            value = (value << 8) + static_cast<uint8_t>(data[i + j]);
        }

        result += chars[(value >> 18) & 0x3F];
        result += chars[(value >> 12) & 0x3F];
        if (i + 1 < data.size()) {
            result += chars[(value >> 6) & 0x3F];
        } else {
            result += '=';
        }
        if (i + 2 < data.size()) {
            result += chars[value & 0x3F];
        } else {
            result += '=';
        }
    }

    return result;
}

// BasicAuthenticator implementation
BasicAuthenticator::BasicAuthenticator(std::string realm) : realm_(std::move(realm)) {}

Result<AuthResult> BasicAuthenticator::authenticate(const Credentials& credentials) {
    if (credentials.username.empty() || credentials.password.empty()) {
        return Error::validation_error("Username and password are required for basic authentication");
    }

    // Create authentication token (in real implementation, this would be validated against a backend)
    std::string auth_string = credentials.username + ":" + credentials.password;
    std::string encoded_credentials = base64_encode(auth_string);

    AuthResult result;
    result.success = true;
    result.access_token = "Basic " + encoded_credentials;
    result.token_type = "Basic";
    result.expires_in = 3600; // 1 hour default

    return result;
}

Result<AuthResult> BasicAuthenticator::refresh(const std::string& refresh_token) {
    return Error::unsupported_operation("Basic authentication does not support token refresh");
}

bool BasicAuthenticator::validate_token(const std::string& access_token) {
    if (access_token.substr(0, 6) != "Basic ") {
        return false;
    }

    // Basic validation - in real implementation, would decode and validate credentials
    try {
        std::string encoded = access_token.substr(6);
        // Check if it looks like base64 encoded
        return !encoded.empty() && encoded.find(" ") == std::string::npos;
    } catch (...) {
        return false;
    }
}

std::string BasicAuthenticator::get_scheme() const {
    return "Basic";
}

std::vector<AuthChallenge> BasicAuthenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "Basic",
        .realm = realm_,
        .parameters = {}
    }};
}

bool BasicAuthenticator::supports_refresh() const {
    return false;
}

// BearerAuthenticator implementation
BearerAuthenticator::BearerAuthenticator(std::string realm) : realm_(std::move(realm)) {}

Result<AuthResult> BearerAuthenticator::authenticate(const Credentials& credentials) {
    if (credentials.bearer_token && !credentials.bearer_token->empty()) {
        AuthResult result;
        result.success = true;
        result.access_token = *credentials.bearer_token;
        result.token_type = "Bearer";
        result.expires_in = 3600; // Default 1 hour

        return result;
    }

    return Error::validation_error("Bearer token is required for bearer authentication");
}

Result<AuthResult> BearerAuthenticator::refresh(const std::string& refresh_token) {
    if (!refresh_token.empty()) {
        AuthResult result;
        result.success = true;
        result.access_token = "new_bearer_token_" + std::to_string(std::time(nullptr));
        result.refresh_token = refresh_token;
        result.token_type = "Bearer";
        result.expires_in = 3600;

        return result;
    }

    return Error::validation_error("Refresh token is required");
}

bool BearerAuthenticator::validate_token(const std::string& access_token) {
    if (access_token.substr(0, 7) != "Bearer ") {
        return false;
    }

    // Basic validation - in real implementation, would validate JWT or check with auth server
    std::string token = access_token.substr(7);
    return !token.empty() && token.length() > 10; // Basic length check
}

std::string BearerAuthenticator::get_scheme() const {
    return "Bearer";
}

std::vector<AuthChallenge> BearerAuthenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "Bearer",
        .realm = realm_,
        .parameters = {}
    }};
}

bool BearerAuthenticator::supports_refresh() const {
    return true;
}

// ApiKeyAuthenticator implementation
ApiKeyAuthenticator::ApiKeyAuthenticator(std::string header_name, std::string realm)
    : header_name_(std::move(header_name)), realm_(std::move(realm)) {}

Result<AuthResult> ApiKeyAuthenticator::authenticate(const Credentials& credentials) {
    if (credentials.api_key && !credentials.api_key->empty()) {
        AuthResult result;
        result.success = true;
        result.access_token = *credentials.api_key;
        result.token_type = "ApiKey";
        result.expires_in = 0; // API keys typically don't expire

        return result;
    }

    return Error::validation_error("API key is required for API key authentication");
}

Result<AuthResult> ApiKeyAuthenticator::refresh(const std::string& refresh_token) {
    return Error::unsupported_operation("API key authentication does not support token refresh");
}

bool ApiKeyAuthenticator::validate_token(const std::string& access_token) {
    // Basic validation - in real implementation, would validate with API management
    return !access_token.empty() && access_token.length() >= 16; // Basic length check
}

std::string ApiKeyAuthenticator::get_scheme() const {
    return "ApiKey";
}

std::vector<AuthChallenge> ApiKeyAuthenticator::get_challenges() const {
    return {AuthChallenge{
        .scheme = "ApiKey",
        .realm = realm_,
        .parameters = {{"header_name", header_name_}}
    }};
}

bool ApiKeyAuthenticator::supports_refresh() const {
    return false;
}

// AuthenticatorFactory implementation
struct AuthenticatorFactory::Impl {
    std::unordered_map<std::string, std::function<std::unique_ptr<Authenticator>()>> creators;
    mutable std::mutex mutex;

    Impl() {
        // Register built-in authenticators
        creators["basic"] = []() { return std::make_unique<BasicAuthenticator>(); };
        creators["bearer"] = []() { return std::make_unique<BearerAuthenticator>(); };
        creators["apikey"] = []() { return std::make_unique<ApiKeyAuthenticator>(); };
    }
};

AuthenticatorFactory::AuthenticatorFactory() : impl_(std::make_unique<Impl>()) {}

AuthenticatorFactory& AuthenticatorFactory::instance() {
    static AuthenticatorFactory instance;
    return instance;
}

void AuthenticatorFactory::register_authenticator(const std::string& scheme,
                                                std::function<std::unique_ptr<Authenticator>()> factory) {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    impl_->creators[scheme] = std::move(factory);
}

std::unique_ptr<Authenticator> AuthenticatorFactory::create_authenticator(const std::string& scheme) {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    auto it = impl_->creators.find(scheme);
    if (it == impl_->creators.end()) {
        return nullptr;
    }
    return it->second();
}

std::vector<std::string> AuthenticatorFactory::get_registered_schemes() const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    std::vector<std::string> schemes;
    for (const auto& [scheme, _] : impl_->creators) {
        schemes.push_back(scheme);
    }
    return schemes;
}

bool AuthenticatorFactory::is_supported(const std::string& scheme) const {
    std::lock_guard<std::mutex> lock(impl_->mutex);
    return impl_->creators.find(scheme) != impl_->creators.end();
}

// Utility functions
Result<AuthResult> authenticate_with_scheme(const std::string& scheme,
                                             const Credentials& credentials,
                                             const std::string& realm) {
    auto& factory = AuthenticatorFactory::instance();
    auto authenticator = factory.create_authenticator(scheme);

    if (!authenticator) {
        return Error::validation_error("Unsupported authentication scheme: " + scheme);
    }

    return authenticator->authenticate(credentials);
}

std::vector<AuthChallenge> get_supported_challenges() {
    auto& factory = AuthenticatorFactory::instance();
    std::vector<AuthChallenge> challenges;

    for (const auto& scheme : factory.get_registered_schemes()) {
        auto authenticator = factory.create_authenticator(scheme);
        if (authenticator) {
            auto scheme_challenges = authenticator->get_challenges();
            challenges.insert(challenges.end(),
                            scheme_challenges.begin(),
                            scheme_challenges.end());
        }
    }

    return challenges;
}

std::string generate_www_authenticate_header(const std::vector<AuthChallenge>& challenges) {
    if (challenges.empty()) {
        return "";
    }

    std::ostringstream ss;
    for (size_t i = 0; i < challenges.size(); ++i) {
        const auto& challenge = challenges[i];

        ss << challenge.scheme;
        if (!challenge.realm.empty()) {
            ss << " realm=\"" << challenge.realm << "\"";
        }

        for (const auto& [key, value] : challenge.parameters) {
            ss << " " << key << "=\"" << value << "\"";
        }

        if (i < challenges.size() - 1) {
            ss << ", ";
        }
    }

    return ss.str();
}

AuthChallenge parse_www_authenticate_header(const std::string& header) {
    AuthChallenge challenge;

    size_t space_pos = header.find(' ');
    if (space_pos == std::string::npos) {
        challenge.scheme = header;
        return challenge;
    }

    challenge.scheme = header.substr(0, space_pos);

    // Parse parameters (simplified implementation)
    std::string params_str = header.substr(space_pos + 1);
    size_t start = 0;

    while (start < params_str.length()) {
        // Skip whitespace
        while (start < params_str.length() && std::isspace(params_str[start])) {
            start++;
        }

        if (start >= params_str.length()) {
            break;
        }

        // Find key
        size_t equal_pos = params_str.find('=', start);
        if (equal_pos == std::string::npos) {
            break;
        }

        std::string key = params_str.substr(start, equal_pos - start);
        start = equal_pos + 1;

        // Skip whitespace
        while (start < params_str.length() && std::isspace(params_str[start])) {
            start++;
        }

        if (start >= params_str.length()) {
            break;
        }

        // Find value (quoted)
        char quote = params_str[start];
        if (quote != '"' && quote != '\'') {
            break;
        }

        start++;
        size_t end_pos = params_str.find(quote, start);
        if (end_pos == std::string::npos) {
            break;
        }

        std::string value = params_str.substr(start, end_pos - start);
        start = end_pos + 1;

        if (key == "realm") {
            challenge.realm = value;
        } else {
            challenge.parameters[key] = value;
        }
    }

    return challenge;
}

} // namespace auth
} // namespace mcp